package com.ruyuan.ingestion.sink;

import com.ruyuan.ingestion.common.Parsable;
import com.ruyuan.ingestion.common.Persistable;
import com.ruyuan.ingestion.config.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ConnectionFactory;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.Table;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

/**
 * 数据写入Hbase
 * */
public class HBaseWriter implements Persistable {

    private String coreSite = null;
    private String hdfsSite = null;
    private String hbaseSite = null;

    private String hbTable = null;
    private Parsable<Put> parser = null;

    public HBaseWriter(String hbTable,Parsable<Put> parser) {
        this.hbTable = hbTable;
        this.parser = parser;
    }

    @Override
    public void initialize(Properties properties) {
        this.coreSite = properties.getProperty(Configuration.CORE_SITE);
        this.hdfsSite = properties.getProperty(Configuration.HDFS_SITE);
        this.hbaseSite = properties.getProperty(Configuration.HBASE_SITE);
    }

    /**
     * 返回值：写入的数据数量
     * */
    @Override
    public int write(ConsumerRecords<String,String> records) throws Exception {
        int numCounts;
        if (this.hbaseSite ==null || this.hbaseSite.isEmpty()) {
            throw new Exception("hbase-site.xml 没有设置");
        }

        org.apache.hadoop.conf.Configuration configuration = HBaseConfiguration.create();
        if (this.coreSite != null) {
            configuration.addResource(new Path(this.coreSite));
        }
        if (this.hdfsSite != null) {
            configuration.addResource(new Path(this.hdfsSite));
        }
        configuration.addResource(new Path(this.hbaseSite));

        Connection connection = ConnectionFactory.createConnection(configuration);
        try {
            Table table = connection.getTable(TableName.valueOf(this.hbTable));

            try {
                //批量的插入到HBase
                List<Put> puts = new ArrayList<>();
                /**
                 * 对于头信息标记判断，只判断一次就可以了
                 * */
                long passHead = 0;
                /**
                 * 遍历records，放入到puts集合
                 * */
                for(ConsumerRecord<String,String> record:records) {
                    try {
                        String[] elements = record.value().split(",", -1);
                        if (passHead ==0 && this.parser.isHeader(elements)) {
                            passHead = 1;
                            continue;
                        }

                        //数据通过check后然后放入到puts集合
                        if (this.parser.isValid(elements)) {
                            puts.add(this.parser.parse(elements));
                        }else {
//                            log.info("invalid message")
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }

                //批量写入
                if (puts.size() > 0) {
                    table.put(puts);
                }

                //记录数
                numCounts = puts.size();
            } finally {
                table.close();
            }

        } finally {
            connection.close();
        }
        return numCounts;
    }
}